package com.aizuda.easy.security.filter.wrapper;

import com.aizuda.easy.security.HandlerFactory;
import com.aizuda.easy.security.handler.FunctionHandler;
import com.aizuda.easy.security.handler.RepFunctionHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;

public class RepSSEWrapper extends HttpServletResponseWrapper {
    private static final Logger log = LoggerFactory.getLogger(RepSSEWrapper.class);

    private final SseServletOutputStream outputStream;
    private PrintWriter writer;
    private HandlerFactory factory;
    private HttpServletResponse response;

    public RepSSEWrapper(HttpServletResponse response, HandlerFactory factory) throws IOException {
        super(response);
        this.response = response;
        this.outputStream = new SseServletOutputStream(response.getOutputStream());
        this.factory = factory;
    }

    @Override
    public ServletOutputStream getOutputStream() {
        return outputStream;
    }

    @Override
    public PrintWriter getWriter() {
        if (writer == null) {
            writer = new PrintWriter(new OutputStreamWriter(outputStream, StandardCharsets.UTF_8));
        }
        return writer;
    }


    // 自定义 OutputStream 处理 SSE 数据
    class SseServletOutputStream extends ServletOutputStream {
        private final ServletOutputStream originalStream;

        public SseServletOutputStream(ServletOutputStream originalStream) {
            this.originalStream = originalStream;
        }

        @Override
        public boolean isReady() {
            return originalStream.isReady();
        }

        @Override
        public void setWriteListener(WriteListener listener) {
            originalStream.setWriteListener(listener);
        }

        @Override
        public void write(int b) throws IOException {
            // 单字节写入（效率较低，实际场景建议重写 write(byte[])）
            originalStream.write(b);
        }

        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            // 拦截并修改 SSE 数据块
            String originalData = new String(b, off, len, StandardCharsets.UTF_8);
            String modifiedData = modifySseData(originalData);
            byte[] modifiedBytes = modifiedData.getBytes(StandardCharsets.UTF_8);
            originalStream.write(modifiedBytes);
        }

        private String modifySseData(String body) throws IOException {
            for (FunctionHandler functionHandler : factory.getFunctionHandlers()) {
                if (functionHandler instanceof RepFunctionHandler) {
                    log.debug("exec handler : {}", functionHandler.getClass().getName());
                    body = ((RepFunctionHandler) functionHandler).exec(response, body);
                }
            }
            return body;
        }

    }
}


